import pyrealsense2 as rs
import numpy as np
import cv2


class D435(object):
    def __init__(self):
        self.pipeline = rs.pipeline()
        self.config = rs.config()
        self.config.enable_stream(rs.stream.depth, 848, 480, rs.format.z16, 30)
        self.config.enable_stream(rs.stream.color, 848, 480, rs.format.bgr8, 30)
        # Get device product line for setting a supporting resolution
        pipeline_wrapper = rs.pipeline_wrapper(self.pipeline)
        pipeline_profile = self.config.resolve(pipeline_wrapper)
        device = pipeline_profile.get_device()
        self.out_param_halcon_type = None

        found_rgb = False
        for s in device.sensors:
            if s.get_info(rs.camera_info.name) == 'RGB Camera':
                found_rgb = True
                break
        if not found_rgb:
            print("The demo requires Depth camera with Color sensor")
            exit(0)

        # Start streaming

        profile = self.pipeline.start(self.config)
        # self.intr = profile.as_video_stream_profile().get_intrinsics()
        # print(intr)
        # Getting the depth sensor's depth scale (see rs-align example for explanation)
        depth_sensor = profile.get_device().first_depth_sensor()
        color_sensor = profile.get_device().first_color_sensor()
        # color_sensor.set_option(rs.option.exposure, 100)
        # color_sensor.set_option(rs.option.gain, 12.00)
        # color_sensor.set_option(rs.option.saturation, 72.00)
        # color_sensor.set_option(rs.option.sharpness, 100.00)
        self.depth_scale = depth_sensor.get_depth_scale()
        # print("Depth Scale is: ", self.depth_scale)

        # We will be removing the background of objects more than
        #  clipping_distance_in_meters meters away
        # clipping_distance_in_meters = 1  # 1 meter
        # clipping_distance = clipping_distance_in_meters / self.depth_scale

        # Create an align object
        # rs.align allows us to perform alignment of depth frames to others frames
        # The "align_to" is the stream type to which we plan to align depth frames.
        align_to = rs.stream.color
        self.align = rs.align(align_to)
        # device_product_line = str(device.get_info(rs.camera_info.product_line))

        frames = self.pipeline.wait_for_frames()

        aligned_frames = self.align.process(frames)

        # Get aligned frames
        depth_frame = aligned_frames.get_depth_frame()  # aligned_depth_frame is a 640x480 depth image
        color_frame = aligned_frames.get_color_frame()
        self.depth_intrinsics = rs.video_stream_profile(
            depth_frame.profile).get_intrinsics()
        self.color_intrinsics = rs.video_stream_profile(
            color_frame.profile).get_intrinsics()
        # depth_frame = frames.get_depth_frame()
        # color_frame = frames.get_color_frame()
        # Convert images to numpy arrays
        self.depth_image = np.asanyarray(depth_frame.get_data())
        self.color_image = np.asanyarray(color_frame.get_data())
        # Apply colormap on depth image (image must be converted to 8-bit per pixel first)
        self.depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(self.depth_image, alpha=0.03), cv2.COLORMAP_JET)

    def update(self):
        frames = self.pipeline.wait_for_frames()
        aligned_frames = self.align.process(frames)

        # Get aligned frames
        depth_frame = aligned_frames.get_depth_frame()  # aligned_depth_frame is a 640x480 depth image
        color_frame = aligned_frames.get_color_frame()
        # Convert images to numpy arrays
        self.depth_image = np.asanyarray(depth_frame.get_data())
        self.color_image = np.asanyarray(color_frame.get_data())
        # Apply colormap on depth image (image must be converted to 8-bit per pixel first)
        self.depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(self.depth_image, alpha=0.03), cv2.COLORMAP_JET)

    def get_images(self):
        return self.color_image, self.depth_image

    def pixel_to_3d(self, u, v, d, depth_scale):
        fx = self.color_intrinsics.fx
        fy = self.color_intrinsics.fy
        u0 = self.color_intrinsics.ppx
        v0 = self.color_intrinsics.ppy
        z = d / depth_scale
        x = (u - u0) * z / fx
        y = (v - v0) * z / fy
        point_3d = np.array([x, y, z])
        return point_3d

    def close(self):
        self.config.disable_stream(rs.stream.depth)
        self.config.disable_stream(rs.stream.color)
        self.pipeline.stop()

    def set_param_outer(self, file_path):
        content = []
        with open(file_path, "r") as f:
            for line in f:
                str_list = line.split()
                for str in str_list:
                    content.append(float(str))
        self.out_param_halcon_type = content


if __name__ == '__main__':
    camera = D435()
    inr = camera.color_intrinsics
    while True:
        camera.update()
        color_image, depth_image = camera.get_images()
        cv2.imshow("detection", color_image)
        cv2.imshow("depth", depth_image)
        cv2.waitKey(30)
